/*************************************************************************
    > File Name: AES.cpp
    > Author: ZhaoLulu@BJTU
    > E-mail: 19281121@bjtu.edu.cn
    > Created Time: 2021年10月29日 星期五
 ************************************************************************/
#include <iostream>
#include <string>
#include <fstream>
#include "AES.h"

using namespace std;

/*******************************密钥扩展函数*******************************/
// 将 4 个 byte 合并为一个 word
word Word(byte &k1, byte &k2, byte &k3, byte &k4) {
    word result(0x00000000);
    word temp;
    temp = k1.to_ulong();
    temp <<= 24;
    result |= temp;
    temp = k2.to_ulong();
    temp <<= 16;
    result |= temp;
    temp = k3.to_ulong();
    temp <<= 8;
    result |= temp;
    temp = k4.to_ulong();
    result |= temp;
    return result;
}

// 按字节循环左移 1 位（[a0, a1, a2, a3]变为[a1, a2, a3, z0]）
word RotWord(word &rw) {
    word high = rw << 8;
    word low = rw >> 24;
    return high | low;
}

// 对输入 word 中每一个 byte 进行 S 盒变换
word SubWord(word sw) {
    word temp;
    for (int i = 0; i < 32; i += 8) {
        int row = sw[i + 7] * 8 + sw[i + 6] * 4 + sw[i + 5] * 2 + sw[i + 4] * 1;
        int col = sw[i + 3] * 8 + sw[i + 2] * 4 + sw[i + 1] * 2 + sw[i + 0] * 1;
        byte val = S_Box[row][col];
        for (int j = 0; j < 8; ++j) {
            temp[i + j] = val[j];
        }
    }
    return temp;
}

// 密钥扩展函数
void KeyExpansion(byte key[4 * Nk], word w[4 * (Nr + 1)]) {
    word temp;
    int i = 0;
    // w[] 前 4 个就是用户输入密钥
    while (i < Nk) {
        w[i] = Word(key[4 * i], key[4 * i + 1], key[4 * i + 2], key[4 * i + 3]);
        i++;
    }
    i = Nk;
    while (i < 4 * (Nr + 1)) {
        // 记录下前一个 word
        temp = w[i - 1];
        // 如果是 Nk 的倍数位置
        if (i % Nk == 0) {
            w[i] = w[i - Nk] ^ SubWord(RotWord(temp)) ^ Rcon[i / Nk - 1];
        } else {
            w[i] = w[i - Nk] ^ temp;
        }
        i++;
    }
}

/*******************************加密变换函数*******************************/
// S 盒变换
void SubBytes(byte mtx[4 * 4]) {
    for (int i = 0; i < 16; ++i) {
        int row = mtx[i][7] * 8 + mtx[i][6] * 4 + mtx[i][5] * 2 + mtx[i][4] * 1;
        int col = mtx[i][3] * 8 + mtx[i][2] * 4 + mtx[i][1] * 2 + mtx[i][0] * 1;
        mtx[i] = S_Box[row][col];
    }
}

// 行变换
void ShiftRows(byte mtx[4 * 4]) {
    // 第 1 行不变
    // 第 2 行循环左移 1 位
    byte temp = mtx[4];
    for (int i = 0; i < 3; ++i) {
        mtx[i + 4] = mtx[i + 5];
    }
    mtx[7] = temp;
    // 第 3 行循环左移 2 位
    for (int i = 0; i < 2; ++i) {
        temp = mtx[i + 8];
        mtx[i + 8] = mtx[i + 10];
        mtx[i + 10] = temp;
    }
    // 第 4 行循环左移 3 位
    temp = mtx[15];
    for (int i = 3; i > 0; --i) {
        mtx[i + 12] = mtx[i + 11];
    }
    mtx[12] = temp;
}

// 有限域上的乘法运算
byte GFMul(byte a, byte b) {
    byte p = 0;
    byte hi_bit_set;
    for (int counter = 0; counter < 8; ++counter) {
        if ((b & byte(1)) != 0) {
            p ^= a;
        }
        hi_bit_set = (byte) (a & byte(0x80));
        a <<= 1;
        if (hi_bit_set != 0) {
            // x^4 + x^3 + x + 1
            a ^= 0x1b;
        }
        b >>= 1;
    }
    return p;
}

// 列变换
/*
 * s[0][c]'   02 03 01 01   s[0][c]
 * s[1][c]' = 01 02 03 01 * s[1][c]  0 <= c < 4
 * s[2][c]'   01 01 02 03   s[2][c]
 * s[3][c]'   03 01 01 02   s[3][c]
 */
void MixColumns(byte mtx[4 * 4]) {
    byte arr[4];
    for (int i = 0; i < 4; ++i) {
        for (int j = 0; j < 4; ++j) {
            arr[j] = mtx[i + j * 4];
        }
        mtx[i] = GFMul(0x02, arr[0]) ^ GFMul(0x03, arr[1]) ^ arr[2] ^ arr[3];
        mtx[i + 4] = arr[0] ^ GFMul(0x02, arr[1]) ^ GFMul(0x03, arr[2]) ^ arr[3];
        mtx[i + 8] = arr[0] ^ arr[1] ^ GFMul(0x02, arr[2]) ^ GFMul(0x03, arr[3]);
        mtx[i + 12] = GFMul(0x03, arr[0]) ^ arr[1] ^ arr[2] ^ GFMul(0x02, arr[3]);
    }
}

// 轮密钥加变换
void AddRoundKey(byte mtx[4 * 4], word k[4]) {
    for (int i = 0; i < 4; ++i) {
        word k1 = k[i] >> 24;
        word k2 = (k[i] << 8) >> 24;
        word k3 = (k[i] << 16) >> 24;
        word k4 = (k[i] << 24) >> 24;
        mtx[i] = mtx[i] ^ byte(k1.to_ulong());
        mtx[i + 4] = mtx[i + 4] ^ byte(k2.to_ulong());
        mtx[i + 8] = mtx[i + 8] ^ byte(k3.to_ulong());
        mtx[i + 12] = mtx[i + 12] ^ byte(k4.to_ulong());
    }
}

/*******************************解密变换函数*******************************/
// 逆 S 盒变换
void InvSubBytes(byte mtx[4 * 4]) {
    for (int i = 0; i < 16; ++i) {
        int row = mtx[i][7] * 8 + mtx[i][6] * 4 + mtx[i][5] * 2 + mtx[i][4] * 1;
        int col = mtx[i][3] * 8 + mtx[i][2] * 4 + mtx[i][1] * 2 + mtx[i][0] * 1;
        mtx[i] = Inv_S_Box[row][col];
    }
}

// 逆行变换
void InvShiftRows(byte mtx[4 * 4]) {
    // 第 1 行不变
    // 第 2 行循环右移 1 位
    byte temp = mtx[7];
    for (int i = 3; i > 0; --i) {
        mtx[i + 4] = mtx[i + 3];
    }
    mtx[4] = temp;
    // 第 3 行循环右移 2 位
    for (int i = 0; i < 2; ++i) {
        temp = mtx[i + 8];
        mtx[i + 8] = mtx[i + 10];
        mtx[i + 10] = temp;
    }
    // 第 4 行循环右移 3 位
    temp = mtx[12];
    for (int i = 0; i < 3; ++i) {
        mtx[i + 12] = mtx[i + 13];
    }
    mtx[15] = temp;
}

// 逆列变换
/*
 * s[0][c]'   0e 0b 0d 09   s[0][c]
 * s[1][c]' = 09 0e 0b 0d * s[1][c]  0 <= c < 4
 * s[2][c]'   0d 09 0e 0b   s[2][c]
 * s[3][c]'   0b 0d 09 0e   s[3][c]
 */
void InvMixColumns(byte mtx[4 * 4]) {
    byte arr[4];
    for (int i = 0; i < 4; ++i) {
        for (int j = 0; j < 4; ++j) {
            arr[j] = mtx[i + j * 4];
        }
        mtx[i] = GFMul(0x0e, arr[0]) ^ GFMul(0x0b, arr[1]) ^ GFMul(0x0d, arr[2]) ^ GFMul(0x09, arr[3]);
        mtx[i + 4] = GFMul(0x09, arr[0]) ^ GFMul(0x0e, arr[1]) ^ GFMul(0x0b, arr[2]) ^ GFMul(0x0d, arr[3]);
        mtx[i + 8] = GFMul(0x0d, arr[0]) ^ GFMul(0x09, arr[1]) ^ GFMul(0x0e, arr[2]) ^ GFMul(0x0b, arr[3]);
        mtx[i + 12] = GFMul(0x0b, arr[0]) ^ GFMul(0x0d, arr[1]) ^ GFMul(0x09, arr[2]) ^ GFMul(0x0e, arr[3]);
    }
}

/*******************************加密解密函数*******************************/
void encrypt(byte in[4 * 4], word w[4 * (Nr + 1)]) {
    word key[4];
    for (int i = 0; i < 4; ++i) {
        key[i] = w[i];
    }
    AddRoundKey(in, key);
    for (int round = 1; round < Nr; ++round) {
        SubBytes(in);
        ShiftRows(in);
        MixColumns(in);
        for (int i = 0; i < 4; ++i) {
            key[i] = w[4 * round + i];
        }
        AddRoundKey(in, key);
    }
    SubBytes(in);
    ShiftRows(in);
    for (int i = 0; i < 4; ++i) {
        key[i] = w[4 * Nr + i];
    }
    AddRoundKey(in, key);
}

void decrypt(byte in[4 * 4], word w[4 * (Nr + 1)]) {
    word key[4];
    for (int i = 0; i < 4; ++i) {
        key[i] = w[4 * Nr + i];
    }
    AddRoundKey(in, key);
    for (int round = Nr - 1; round > 0; --round) {
        InvShiftRows(in);
        InvSubBytes(in);
        for (int i = 0; i < 4; ++i) {
            key[i] = w[4 * round + i];
        }
        AddRoundKey(in, key);
        InvMixColumns(in);
    }
    InvShiftRows(in);
    InvSubBytes(in);
    for (int i = 0; i < 4; ++i) {
        key[i] = w[i];
    }
    AddRoundKey(in, key);
}

/*******************************工作模式函数*******************************/
// 将 char 字符数组转换为二进制存入 byte 数组
void charToByte(byte out[16], const char s[16]) {
    for (int i = 0; i < 16; ++i) {
        for (int j = 0; j < 8; ++j) {
            out[i][j] = ((s[i] >> j) & 1);
        }
    }
}

// 将连续的 128 位分成 16 组，存入 byte 数组
void divideToByte(byte out[16], bitset<128> &data) {
    bitset<128> temp;
    for (int i = 0; i < 16; ++i) {
        temp = (data << 8 * i) >> 120;
        out[i] = temp.to_ulong();
    }
}

// 将 16 个 byte 合并成连续 128 位
bitset<128> mergeByte(byte in[16]) {
    bitset<128> res;
    res.reset();
    bitset<128> temp;
    for (int i = 0; i < 16; ++i) {
        temp = in[i].to_ulong();
        temp <<= 8 * (15 - i);
        res |= temp;
    }
    return res;
}

/*******************************样例测试函数*******************************/
// 128 位明文样例测试
void CaseTest1() {
    for (int i = 0; i < 10; ++i) {
        // 构造文件路径
        string path;
        path.append("case");
        path.append(to_string(i + 1));
        path.append(".txt");

        // 读取文件
        ifstream in;
        in.open(path, ios::binary);

        // 获取密钥
        string keyStr;
        byte key[16];
        getline(in, keyStr);
        charToByte(key, keyStr.c_str());

        // 密钥扩展
        word w[4 * (Nr + 1)];
        KeyExpansion(key, w);

        // 获取明文
        string plaintextStr;
        byte plaintext[16];
        getline(in, plaintextStr);
        in.close();

        // 加密
        charToByte(plaintext, plaintextStr.c_str());
        encrypt(plaintext, w);

        // 写入文件
        ofstream out;
        out.open(path, ios::binary | ios::app);
        for (int j = 0; j < 15; ++j) {
            out << hex << plaintext[j].to_ulong() << " ";
        }
        out << hex << plaintext[15].to_ulong() << endl;
        out.close();
    }
}

// 文件加密测试
void CaseTest2() {
    // 密钥类型转换
    string keyStr = "abcdefghijklmnop"; // 128 位
    byte key[16];
    charToByte(key, keyStr.c_str());

    // 密钥扩展
    word w[4 * (Nr + 1)];
    KeyExpansion(key, w);

    // 加密
    bitset<128> data;
    byte plain[16];
    ifstream in;
    ofstream out;
    in.open("table.png", ios::binary);
    out.open("cipher.txt", ios::binary);
    while (in.read((char *) &data, sizeof data)) {
        divideToByte(plain, data);
        encrypt(plain, w);
        data = mergeByte(plain);
        out.write((char *) &data, sizeof(data));
        data.reset(); // 置0
    }
    in.close();
    out.close();

    // 解密
    in.open("cipher.txt", ios::binary);
    out.open("table_decrypted.jpg", ios::binary);
    while (in.read((char *) &data, sizeof(data))) {
        divideToByte(plain, data);
        decrypt(plain, w);
        data = mergeByte(plain);
        out.write((char *) &data, sizeof(data));
        data.reset(); // 置0
    }
    in.close();
    out.close();
}


int main() {
    // 128 位明文样例测试
    // CaseTest1();

    // 文件加密测试
    // CaseTest2();

    return 0;
}
